package com.lnn.springBase.utils.image;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;

public class ImageUtil {

    /**
     * 将 double[][] 数组转换为灰度图像，NaN 值设为透明。
     * 自动计算数据的最大值和最小值进行归一化。
     * @param data 输入的二维数组
     * @return 生成的 BufferedImage 图像
     */
    public static BufferedImage convertToImage(double[][] data) {
        // 计算最大值和最小值
        double min = Double.POSITIVE_INFINITY;
        double max = Double.NEGATIVE_INFINITY;
        boolean hasValid = false;

        for (double[] row : data) {
            if (row == null) continue;
            for (double value : row) {
                if (!Double.isNaN(value)) {
                    hasValid = true;
                    if (value < min) min = value;
                    if (value > max) max = value;
                }
            }
        }
        if (!hasValid) {
            min = 0;
            max = 0;
        }
        return convertToImage(data, min, max);
    }

    /**
     * 将 double[][] 数组转换为灰度图像，NaN 值设为透明。
     * 使用指定的最大值和最小值进行归一化。
     * @param data 输入的二维数组
     * @param min  数据的最小值（用于归一化）
     * @param max  数据的最大值（用于归一化）
     * @return 生成的 BufferedImage 图像
     */
    public static BufferedImage convertToImage(double[][] data, double min, double max) {
        if (data == null || data.length == 0) {
            return new BufferedImage(0, 0, BufferedImage.TYPE_INT_ARGB);
        }
        int width = data.length;
        int height = data[0].length;
        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);
        for (int x = 0; x < width; x++) {
            if (data[x] == null || data[x].length != height) {
                throw new IllegalArgumentException("数据数组维度不一致");
            }
            for (int y = 0; y < height; y++) {
                double value = data[x][y];
                if (Double.isNaN(value)) {
                    // 设置为完全透明
                    image.setRGB(x, y, 0x00000000);
                } else {
                    // 归一化到 [0, 255]
                    double normalized = (value - min) / (max - min);
                    normalized = Math.max(0.0, Math.min(1.0, normalized)); // 确保在范围内
                    int gray = (int) (normalized * 255.0);
                    // 创建 ARGB 像素值（Alpha=255 表示不透明）
                    int argb = (0xFF << 24) | (gray << 16) | (gray << 8) | gray;
                    image.setRGB(x, y, argb);
                }
            }
        }
        return image;
    }

    /**
     * 将灰度图像转换回 double[][] 数组，透明像素恢复为 NaN。
     * @param image 输入的灰度图像（必须为 ARGB 类型）
     * @param min   原始数据的最小值（用于反归一化）
     * @param max   原始数据的最大值（用于反归一化）
     * @return 还原后的 double[][] 数组
     */
    public static double[][] convertToData(BufferedImage image, double min, double max) {
        if (image == null) {
            return new double[0][0];
        }
        int width = image.getWidth();
        int height = image.getHeight();
        double[][] data = new double[width][height];
        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {
                int argb = image.getRGB(x, y);
                int alpha = (argb >> 24) & 0xFF;

                if (alpha == 0) {
                    data[x][y] = Double.NaN;
                } else {
                    // 提取红色分量（灰度图像 RGB 分量相同）
                    int gray = (argb >> 16) & 0xFF;
                    // 反归一化到原始范围
                    double value = min + (gray / 255.0) * (max - min);
                    data[x][y] = value;
                }
            }
        }
        return data;
    }

    public static void main(String[] args) throws Exception {
        // 创建示例数据（包含 NaN）
        double[][] data = {
                {0.0, 0.5, Double.NaN},
                {0.3, 0.8, 1.0}
        };
        // 转换为图像（自动计算范围）
        BufferedImage image = ImageUtil.convertToImage(data);
        ImageIO.write(image, "PNG", new File("output.png"));
        // 转换回数据（假设我们知道原始范围）
        double min = 0.0;
        double max = 1.0;
        double[][] restoredData = ImageUtil.convertToData(image, min, max);

        // 验证数据还原正确性
        for (int i = 0; i < data.length; i++) {
            for (int j = 0; j < data[i].length; j++) {
                System.out.printf("%6.2f ", restoredData[i][j]);
            }
            System.out.println();
        }
    }
}
